package redislib

import (
	"context"
	"fmt"
	"github.com/go-redis/redis/v8"
	"time"
)

type RedisStream struct {
	client       *redis.Client
	ctx          context.Context
	streamKey    string
	groupName    string
	consumerName string
}

// 修改 NewRedisStream 函数中的判断逻辑
func NewRedisStream(client *redis.Client, streamKey, groupName, consumerName string) (*RedisStream, error) {
	rs := &RedisStream{
		client:       client,
		ctx:          context.Background(),
		streamKey:    streamKey,
		groupName:    groupName,
		consumerName: consumerName,
	}

	// 使用 XInfoGroups 判断流是否存在（更可靠）
	info, err := rs.client.XInfoStream(rs.ctx, streamKey).Result()
	if err != nil || info == nil {
		// 流不存在，创建并初始化消费者组
		err = rs.client.XGroupCreateMkStream(rs.ctx, streamKey, groupName, "0").Err()
		if err != nil && err.Error() != "BUSYGROUP Consumer Group name already exists" {
			return nil, err
		}
	}
	return rs, nil
}

// Publish 发送消息到 Stream
func (rs *RedisStream) Publish(message map[string]interface{}) error {
	err := rs.client.XAdd(rs.ctx, &redis.XAddArgs{
		Stream: rs.streamKey,
		Values: message,
	}).Err()
	if err != nil {
		fmt.Printf("Failed to publish message: %v\n", err)
	}
	return err
}

// Consume 消费消息并处理
func (rs *RedisStream) Consume(handler func(message map[string]interface{}) error) {
	retryCount := make(map[string]int) // 记录每条消息的重试次数

	for {
		cmd := rs.client.XReadGroup(rs.ctx, &redis.XReadGroupArgs{
			Group:    rs.groupName,
			Consumer: rs.consumerName,
			Streams:  []string{rs.streamKey, ">"},
			Count:    1,
			Block:    5000 * time.Millisecond,
		})

		msgs, err := cmd.Result()
		if err != nil || len(msgs) == 0 {
			break
		}

		for _, msg := range msgs[0].Messages {
			msgID := msg.ID

			// 执行带重试的处理逻辑
			err := retryOnFailure(func() error {
				return handler(msg.Values)
			}, 3)

			if err != nil {
				fmt.Printf("Message ID: %s failed after 3 retries. Giving up.\n", msgID)
			} else {
				// 只有成功才确认消息
				rs.client.XAck(rs.ctx, rs.streamKey, rs.groupName, msgID)
			}

			// 清除该消息的重试计数
			delete(retryCount, msgID)
		}
	}
}

// retryOnFailure 尝试执行 fn 最多 maxRetries 次
func retryOnFailure(fn func() error, maxRetries int) error {
	var err error
	for i := 0; i < maxRetries; i++ {
		err = fn()
		if err == nil {
			return nil
		}
		fmt.Printf("Attempt %d failed: %v. Retrying...\n", i+1, err)
		time.Sleep(2 * time.Second) // backoff
	}
	return err
}
